package cn.bugstack.ai.service;

import cn.bugstack.ai.model.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;

/**
 * AI 服务
 * 处理 AI 模型调用和流式响应
 * 
 * @author bugstack
 */
@Slf4j
@Service
public class AIService {
    
    private final ChatClient chatClient;
    private final TracingService tracingService;
    private final SessionManager sessionManager;
    
    public AIService(ChatModel chatModel, TracingService tracingService, SessionManager sessionManager) {
        this.chatClient = ChatClient.builder(chatModel).build();
        this.tracingService = tracingService;
        this.sessionManager = sessionManager;
        log.info("✅ AIService initialized");
    }
    
    /**
     * 流式聊天（普通模式）
     */
    public Flux<AIResponse> streamChat(AIRequest request) {
        String sessionId = request.getSessionId();
        String traceId = tracingService.createTraceId();
        
        log.info("🚀 Starting stream chat - SessionId: {}, TraceId: {}", sessionId, traceId);
        
        // 记录用户输入
        sessionManager.addMessage(sessionId, "User: " + request.getContent());
        
        return chatClient.prompt()
                .user(request.getContent())
                .stream()
                .content()
                .map(chunk -> {
                    log.trace("📨 Streaming chunk - SessionId: {}, Length: {}", sessionId, chunk.length());
                    return AIResponse.message(sessionId, chunk);
                })
                .concatWith(Flux.just(
                        AIResponse.complete(sessionId, traceId)
                ))
                .doOnError(error -> {
                    log.error("❌ Error in stream chat - SessionId: {}, Error: {}", 
                            sessionId, error.getMessage(), error);
                })
                .doOnComplete(() -> {
                    log.info("✅ Stream chat completed - SessionId: {}", sessionId);
                });
    }
    
    /**
     * 流式聊天（溯源模式）
     */
    public Flux<AIResponse> chatWithTracing(AIRequest request) {
        String sessionId = request.getSessionId();
        String traceId = tracingService.createTraceId();
        
        log.info("🚀 Starting stream chat with tracing - SessionId: {}, TraceId: {}", 
                sessionId, traceId);
        
        long startTime = System.currentTimeMillis();
        
        // 步骤1：记录用户输入
        tracingService.recordStep(traceId, 1, "user_input", "用户输入", 
                request.getContent(), null, 0L);
        
        // 发送步骤1信息
        Flux<AIResponse> step1 = Flux.just(
                AIResponse.builder()
                        .type("trace")
                        .sessionId(sessionId)
                        .traceId(traceId)
                        .content("步骤1: 接收用户输入")
                        .metadata(Map.of("step", 1, "action", "user_input"))
                        .build()
        );
        
        // 步骤2：生成第一次响应
        return step1.concatWith(
                chatClient.prompt()
                        .user(request.getContent())
                        .stream()
                        .content()
                        .collectList()
                        .flatMapMany(chunks -> {
                            String firstResponse = String.join("", chunks);
                            long duration = System.currentTimeMillis() - startTime;
                            
                            // 记录第一次生成
                            tracingService.recordStep(traceId, 2, "generate_response", 
                                    "生成AI响应", request.getContent(), firstResponse, duration);
                            
                            log.info("✅ First response generated - SessionId: {}, Length: {}, Duration: {}ms", 
                                    sessionId, firstResponse.length(), duration);
                            
                            // 发送第一次响应
                            Flux<AIResponse> firstResponseFlux = Flux.just(
                                    AIResponse.builder()
                                            .type("trace")
                                            .sessionId(sessionId)
                                            .traceId(traceId)
                                            .content("步骤2: AI 生成响应\n\n" + firstResponse)
                                            .metadata(Map.of("step", 2, "action", "generate_response", 
                                                    "duration", duration))
                                            .build()
                            );
                            
                            // 步骤3：分析生成过程（溯源分析）
                            String tracePrompt = String.format(
                                    "请简要分析以下对话的生成过程（2-3句话即可）：\n" +
                                    "用户问题：%s\n" +
                                    "AI回答：%s\n" +
                                    "分析要点：回答的思路、主要内容和适用性",
                                    request.getContent(), 
                                    firstResponse.length() > 200 ? firstResponse.substring(0, 200) + "..." : firstResponse
                            );
                            
                            long traceStartTime = System.currentTimeMillis();
                            
                            return firstResponseFlux.concatWith(
                                    chatClient.prompt()
                                            .user(tracePrompt)
                                            .stream()
                                            .content()
                                            .collectList()
                                            .flatMapMany(traceChunks -> {
                                                String traceAnalysis = String.join("", traceChunks);
                                                long traceDuration = System.currentTimeMillis() - traceStartTime;
                                                
                                                // 记录溯源分析
                                                tracingService.recordStep(traceId, 3, "trace_analysis", 
                                                        "溯源分析", tracePrompt, traceAnalysis, traceDuration);
                                                
                                                log.info("✅ Trace analysis completed - SessionId: {}, Duration: {}ms", 
                                                        sessionId, traceDuration);
                                                
                                                // 发送溯源分析结果和完整的溯源记录
                                                return Flux.just(
                                                        AIResponse.builder()
                                                                .type("trace")
                                                                .sessionId(sessionId)
                                                                .traceId(traceId)
                                                                .content("步骤3: 生成过程分析\n\n" + traceAnalysis)
                                                                .metadata(Map.of("step", 3, "action", "trace_analysis", 
                                                                        "duration", traceDuration))
                                                                .build(),
                                                        AIResponse.builder()
                                                                .type("trace_complete")
                                                                .sessionId(sessionId)
                                                                .traceId(traceId)
                                                                .finished(true)
                                                                .metadata(Map.of(
                                                                        "traceRecords", tracingService.getTraceRecords(traceId),
                                                                        "summary", tracingService.getTraceSummary(traceId)
                                                                ))
                                                                .build()
                                                );
                                            })
                            );
                        })
        ).doOnError(error -> {
            log.error("❌ Error in trace chat - SessionId: {}, Error: {}", 
                    sessionId, error.getMessage(), error);
        }).doOnComplete(() -> {
            log.info("✅ Trace chat completed - SessionId: {}, TraceId: {}", sessionId, traceId);
        });
    }
}
